# -*- coding: utf-8 -*-
"""
Created on Thu Jun 10 22:24:34 2021

@author: zhuo 木鸟
"""

import pandas as pd
import numpy as np
from statsmodels.tsa.ar_model import AR
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
import matplotlib.pyplot as plt
from series_to_supervised import series_to_supervised

plt.rcParams['font.sans-serif']=['SimHei']    #画图时使用中文字体
plt.rcParams['axes.unicode_minus'] = False


def plot_result(time_range, test, predictions, autoregression=True):
    # plot results
    fig = plt.figure()
    plt.plot(time_range, test, label='Raw Dataset')
    plt.plot(time_range, predictions, color='r', label='Prediction Dataset', linestyle='--')
    plt.xlabel('日期', fontsize=16)
    plt.ylabel('板块指数', fontsize=16)
    plt.title('实际数据和预测数据', fontsize=24)
    plt.grid()
    plt.legend()
    if autoregression:
        fig.savefig('../图片/实际数据和预测数据_自回归.png')
    else:
        fig.savefig('../图片/实际数据和预测数据_naive.png')
    plt.show()
    

def evalute_naive(train, test):
    # 包括模型训练，模型评价和画图代码
    def persistence_model(x):
        return x
    # 构建 persistence 模型，以及预测数据
    y_test_pred = test.shift(1)
    y_test_pred.iloc[0] = train.iloc[-1]
    time_range = test.index
    
    y_train_pred = train.shift(1).iloc[1:]
    
    
    score_in_test = mean_squared_error(test, y_test_pred)
    score_in_train = mean_squared_error(train.iloc[1:], y_train_pred)
    print('Persisence 模型 MSE（训练集）：', score_in_train)
    print('Persisence 模型 MSE（测试集）：', score_in_test)
    # 画图：
    plot_result(time_range, test, y_test_pred, autoregression=False)
    
    time_range = y_train_pred.index
    plot_result(time_range, train.iloc[1:], y_train_pred, autoregression=False)
    

def evalute_autoregression(train, test, n_in=2, verbose=True):
    # 使用自相关预测
    X_train, y_train = train.iloc[:, 0:n_in], train.iloc[:, -1]
    X_test, y_test = test.iloc[:, 0:n_in], test.iloc[:, -1]
    
    lr = LinearRegression()
    lr.fit(X_train, y_train)
    y_test_pred = lr.predict(X_test)
    y_train_pred = lr.predict(X_train)

    score_in_test = mean_squared_error(y_test, y_test_pred)
    score_in_train = mean_squared_error(y_train, y_train_pred)
    print('自相关模型 模型 MSE（训练集）：', score_in_train)
    print('自相关模型 模型 MSE（测试集）：', score_in_test)
    
    time_range = test.index
    if verbose:
        plot_result(time_range, y_test, y_test_pred, autoregression=True)
    
    time_range = train.index
    if verbose:
        plot_result(time_range, y_train, y_train_pred, autoregression=True)
        
    
    print('模型的参数为： ', [lr.intercept_, lr.coef_])
    return lr


if __name__ == '__main__':
    path = r'../附件/中间数据/光伏-建筑板块市值.xlsx'
    series = pd.read_excel(path, header=0, index_col=0, parse_dates=True, squeeze=True)
    
    train, test = series['2019-4-1':'2021-4-30'], series['2021-5-1':'2021-5-28']
    
    evalute_naive(train, test)
    n_in = 2
    series_window = series_to_supervised(series, n_in=n_in, n_out=1, dropnan=True)
    # 拆分数据集
    train, test = series_window['2019-4-1':'2021-4-30'], series_window['2021-5-1':'2021-5-28']
    
    lr = evalute_autoregression(train, test, n_in)
    